import os
import random

import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import transforms

data_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])


class CartoonDataset(Dataset):
    def __init__(self, cartoon_data_path, photo_data_path):
        super(CartoonDataset, self).__init__()

        self.scenery_cartoon_data_path = "F:/cvpr/scenery_cartoon/hosoda"
        self.scenery_photo_data_path = "F:/cvpr/scenery_photo"

        self.face_cartoon_data_path = "F:/cvpr/face_cartoon/pa_face"
        self.face_photo_data_path = "F:/cvpr/face_photo"

        self.scenery_cartoon_images = self.getimages(self.scenery_cartoon_data_path)
        self.scenery_photo_images = self.getimages(self.scenery_photo_data_path)

        self.face_cartoon_images = self.getimages(self.face_cartoon_data_path)
        self.face_photo_images = self.getimages(self.face_photo_data_path)

    def __getitem__(self, index):
        n = min(len(self.face_photo_images),len(self.scenery_photo_images),len(self.scenery_cartoon_images),len(self.face_cartoon_images))
        index = index % n
        # return face
        if random.random()<=0.25:
            cartoon_image=Image.open(self.face_cartoon_data_path+'/'+self.face_cartoon_images[index])
            cartoon_image=data_transform(cartoon_image)

            photo_image=Image.open(self.face_photo_data_path+'/'+self.face_photo_images[index])
            photo_image=data_transform(photo_image)

            return cartoon_image,photo_image
        # return scenery
        else:

            cartoon_image = Image.open(self.scenery_cartoon_data_path + '/' + self.scenery_cartoon_images[index])
            cartoon_image = data_transform(cartoon_image)

            photo_image = Image.open(self.scenery_photo_data_path + '/' + self.scenery_photo_images[index])
            photo_image = data_transform(photo_image)

            return cartoon_image, photo_image

    def __len__(self):
        return min(len(self.face_photo_images),len(self.scenery_photo_images),len(self.scenery_cartoon_images),len(self.face_cartoon_images))

    def getimages(self, data_path):
        data = []
        for image in os.listdir(data_path):
            data.append(image)
        return data

# face_list=

def next_batch(face=True):
    if face:
        idx=np.arange(0,len())
